import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import { z } from 'zod';
import { DatabaseManager } from './database-manager.js';
import * as path from 'path';

// McpToolState接口定义
interface McpToolState {
  input: any;
  context?: any;
}

// 数据库连接信息接口
interface DatabaseConnectionInfo {
  name: string;
  dialect: string;
  host: string;
  port?: number;
  username?: string;
  password?: string;
  database?: string;
  storage?: string;
  pool?: {
    max?: number;
    min?: number;
    idle?: number;
    acquire?: number;
  };
}

// 添加进程错误处理
process.on('uncaughtException', (error) => {
  console.error('Uncaught Exception:', error);
});

process.on('unhandledRejection', (reason, promise) => {
  console.error('Unhandled Rejection at:', promise, 'reason:', reason);
});

process.on('SIGINT', () => {
  console.log('Received SIGINT. Graceful shutdown...');
  process.exit(0);
});

process.on('SIGTERM', () => {
  console.log('Received SIGTERM. Graceful shutdown...');
  process.exit(0);
});

// 数据库类型
export enum DatabaseType {
  RELATIONAL = 'relational' // 关系型数据库(MySQL, PostgreSQL等)
}

export class DatabaseMcpServer {
  private server: any;
  private dbManager: DatabaseManager;
  private transport?: any;

  constructor() {
    this.dbManager = new DatabaseManager();
    
    // 所有服务统一使用database作为名称
    this.server = new McpServer({
      name: 'database',
      version: '1.0.11',
      capabilities: {
        resources: {},
        tools: {},
      },
    });

    this.initializeTools();
  }

  private initializeTools() {
    // 列出所有可用数据库连接
    this.server.tool('listAllDatabases', '列出所有可用的数据库连接及其详细信息。', {
      input: z.object({}).optional()
    }, async () => {
      try {
        // 获取已连接的关系型数据库
        const relationalDbDetails = this.dbManager.getConnectionDetails();
        
        // 为每个数据库添加类型标识
        const relationalDbs = relationalDbDetails.map(db => ({
          ...db,
          type: DatabaseType.RELATIONAL,
          connected: true
        }));
        
        // 所有数据库
        const allDatabases = [...relationalDbs];
        
        const response = {
          success: true,
          data: allDatabases,
          metadata: {
            total: allDatabases.length,
            relational: relationalDbs.length
          }
        };

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(response, null, 2)
            }
          ]
        };
      } catch (error: any) {
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({
                success: false,
                error: error.message,
                details: '获取数据库列表失败。请检查数据库配置。'
              }, null, 2)
            }
          ]
        };
      }
    });
    
    // 通用的表结构描述工具
    this.server.tool('describeTable', '获取表的结构信息，包括字段名称、类型、是否允许空值等。', {
      input: z.object({
        dbName: z.string().describe('数据库连接名称，必须是已配置的数据库之一。使用 listAllDatabases 工具可以查看所有可用的数据库连接。'),
        tableName: z.string().describe('要查询结构的表名。表名区分大小写，取决于数据库设置。')
      })
    }, async (state: McpToolState) => {
      try {
        const { dbName, tableName } = state.input;
        
        const tableInfo = await this.dbManager.describeTable(dbName, tableName);
        const dbType = DatabaseType.RELATIONAL;

        const response = {
          success: true,
          data: tableInfo,
          metadata: {
            dbName,
            tableName,
            dbType
          }
        };

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(response, null, 2)
            }
          ]
        };
      } catch (error: any) {
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({
                success: false,
                error: error.message,
                details: '无法获取表结构。请确认数据库名称和表名是否正确，以及表是否存在。'
              }, null, 2)
            }
          ]
        };
      }
    });
    
    // 通用列出表工具
    this.server.tool('listTables', '列出指定数据库中的所有表。', {
      input: z.object({
        dbName: z.string().describe('数据库连接名称，必须是已配置的数据库之一。使用 listAllDatabases 工具可以查看所有可用的数据库连接。')
      })
    }, async (state: McpToolState) => {
      try {
        const { dbName } = state.input;
        
        // 获取表列表
        const tables = await this.dbManager.listTables(dbName);
        
        const response = {
          success: true,
          data: tables,
          metadata: {
            dbName,
            count: tables.length
          }
        };

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(response, null, 2)
            }
          ]
        };
      } catch (error: any) {
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({
                success: false,
                error: error.message,
                details: '无法获取表列表。请确认数据库名称是否正确。'
              }, null, 2)
            }
          ]
        };
      }
    });
    
    this.initializeRelationalDbTools();
  }

  private initializeRelationalDbTools() {
    // 执行SQL查询工具 - 关系型数据库
    this.server.tool('query', '执行 SQL 查询。用于在指定关系型数据库(MySQL/PostgreSQL等)上运行 SQL 语句，支持参数化查询。参数化查询请使用 ? 作为占位符，并在 params 数组中按顺序提供参数值。', {
      input: z.object({
        dbName: z.string().describe('数据库连接名称，必须是已配置的数据库之一。使用 listAllDatabases 工具可以查看所有可用的数据库连接。'),
        sql: z.string().describe('要执行的 SQL 查询语句。对于参数化查询，使用 ? 作为占位符。例如: "SELECT * FROM users WHERE id = ?"'),
        params: z.array(z.any()).optional().describe('可选的参数数组，用于参数化查询。数组元素的顺序对应 SQL 中 ? 占位符的顺序。例如: [1] 或 ["value", 123]')
      })
    }, async (state: McpToolState) => {
      try {
        const { dbName, sql, params = [] } = state.input;
        const results = await this.dbManager.executeQuery(dbName, sql, params);

        return {
          content: [{
            type: 'text',
            text: JSON.stringify({
              success: true,
              data: results,
              metadata: {
                dbName,
                rowCount: results.length,
                fields: results.length > 0 ? Object.keys(results[0]) : []
              }
            }, null, 2)
          }]
        };
      } catch (error: any) {
        return {
          content: [{
            type: 'text',
            text: JSON.stringify({
              success: false,
              error: error.message,
              details: '查询执行失败。请检查数据库名称是否正确，SQL 语法是否有误，或参数格式是否有问题。'
            }, null, 2)
          }]
        };
      }
    });

    // 添加数据库连接
    this.server.tool('addDatabaseConnection', '添加新的关系型数据库连接。连接信息将被保存并可在后续会话中使用。', {
      input: z.object({
        name: z.string().describe('数据库连接的唯一名称，用于后续引用该连接。'),
        dialect: z.enum(['mysql', 'postgres', 'sqlite', 'mssql', 'mariadb']).describe('数据库类型，支持mysql、postgres、sqlite、mssql和mariadb。'),
        host: z.string().optional().describe('数据库主机地址，默认为localhost。对于sqlite，可忽略此项。'),
        port: z.number().optional().describe('数据库端口号，如MySQL默认为3306，PostgreSQL默认为5432等。'),
        username: z.string().optional().describe('数据库用户名。对于sqlite，可忽略此项。'),
        password: z.string().optional().describe('数据库密码。对于sqlite，可忽略此项。'),
        database: z.string().optional().describe('要连接的数据库名称。对于某些数据库类型是必需的。'),
        storage: z.string().optional().describe('SQLite数据库文件路径。仅当dialect为sqlite时使用。')
      })
    }, async (state: McpToolState) => {
      try {
        const inputParams = state.input;
        
        // 构建符合DatabaseConnectionInfo接口的对象
        const connInfo: DatabaseConnectionInfo = {
          name: inputParams.name,
          dialect: inputParams.dialect,
          host: inputParams.host || 'localhost', // 默认使用localhost
          port: inputParams.port,
          username: inputParams.username,
          password: inputParams.password,
          database: inputParams.database,
          storage: inputParams.storage
        };
        
        // 如果是sqlite且没有指定storage，设置默认路径
        if (inputParams.dialect === 'sqlite' && !inputParams.storage) {
          connInfo.storage = path.join(this.dbManager.getDataWiseDir(), `${inputParams.name}.db`);
        }
        
        const success = await this.dbManager.addDatabaseConnection(connInfo);
        
        const response = {
          success: true,
          data: {
            name: connInfo.name,
            dialect: connInfo.dialect,
            connected: success
          },
          message: success 
            ? `数据库连接 '${connInfo.name}' 已成功添加并连接` 
            : `数据库连接 '${connInfo.name}' 已添加但连接失败，请检查连接信息`
        };

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(response, null, 2)
            }
          ]
        };
      } catch (error: any) {
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({
                success: false,
                error: error.message,
                details: '添加数据库连接失败。请检查连接名称是否已存在，或者连接参数是否有效。'
              }, null, 2)
            }
          ]
        };
      }
    });

    // 测试数据库连接
    this.server.tool('testDatabaseConnection', '测试数据库连接配置是否有效。该工具可用于测试现有连接或验证新连接配置而不保存。', {
      input: z.object({
        mode: z.enum(['existing', 'new']).describe('测试模式：existing表示测试已保存的连接，new表示测试新的连接配置。'),
        name: z.string().describe('对于existing模式，指定要测试的现有连接名称；对于new模式，指定新连接的名称。'),
        dialect: z.enum(['mysql', 'postgres', 'sqlite', 'mssql', 'mariadb']).optional().describe('数据库类型，仅new模式需要。'),
        host: z.string().optional().describe('数据库主机地址，仅new模式需要。'),
        port: z.number().optional().describe('数据库端口号，仅new模式需要。'),
        username: z.string().optional().describe('数据库用户名，仅new模式需要。'),
        password: z.string().optional().describe('数据库密码，仅new模式需要。'),
        database: z.string().optional().describe('数据库名称，仅new模式需要。'),
        storage: z.string().optional().describe('SQLite数据库文件路径，仅当dialect为sqlite时需要。')
      })
    }, async (state: McpToolState) => {
      try {
        const { mode, name, ...connParams } = state.input;
        let success = false;
        
        if (mode === 'existing') {
          // 测试现有连接
          success = await this.dbManager.testConnection(name);
        } else {
          // 测试新连接配置
          if (!connParams.dialect) {
            throw new Error('在new模式下，必须指定数据库类型(dialect)');
          }
          
          // 针对不同数据库类型处理必需参数
          const connInfo: DatabaseConnectionInfo = {
            name,
            dialect: connParams.dialect,
            host: 'localhost', // 默认值
            port: connParams.port,
            username: connParams.username,
            password: connParams.password,
            database: connParams.database,
            storage: connParams.storage
          };
          
          // 根据数据库类型覆盖host
          if (connParams.host) {
            connInfo.host = connParams.host;
          }
          
          // 对于sqlite，设置默认storage
          if (connParams.dialect === 'sqlite' && !connParams.storage) {
            connInfo.storage = path.join(this.dbManager.getDataWiseDir(), `${name}.db`);
          }
          
          success = await this.dbManager.testConnectionConfig(connInfo);
        }
        
        const response = {
          success: true,
          data: {
            name,
            mode,
            connectionTestResult: success
          },
          message: success 
            ? `数据库连接测试成功` 
            : `数据库连接测试失败，请检查连接信息`
        };

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(response, null, 2)
            }
          ]
        };
      } catch (error: any) {
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({
                success: false,
                error: error.message,
                details: '测试数据库连接失败。请检查连接参数是否有效。'
              }, null, 2)
            }
          ]
        };
      }
    });

    // 删除数据库连接
    this.server.tool('deleteDatabaseConnection', '删除指定的关系型数据库连接。该操作会关闭连接并从配置中永久删除。', {
      input: z.object({
        name: z.string().describe('要删除的数据库连接名称。必须是已配置的连接之一。')
      })
    }, async (state: McpToolState) => {
      try {
        const { name } = state.input;
        const success = await this.dbManager.deleteDatabaseConnection(name);
        
        const response = {
          success: true,
          data: {
            name,
            deleted: success
          },
          message: success 
            ? `数据库连接 '${name}' 已成功删除` 
            : `数据库连接 '${name}' 删除操作未完成`
        };

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(response, null, 2)
            }
          ]
        };
      } catch (error: any) {
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({
                success: false,
                error: error.message,
                details: '删除数据库连接失败。请检查连接名称是否正确，或者连接是否存在。'
              }, null, 2)
            }
          ]
        };
      }
    });

    // 更新数据库连接
    this.server.tool('updateDatabaseConnection', '更新已有的关系型数据库连接信息。该操作会关闭现有连接并用新配置重新连接。', {
      input: z.object({
        name: z.string().describe('要更新的数据库连接名称。必须是已配置的连接之一。'),
        dialect: z.enum(['mysql', 'postgres', 'sqlite', 'mssql', 'mariadb']).optional().describe('可选，更新数据库类型。'),
        host: z.string().optional().describe('可选，更新数据库主机地址。'),
        port: z.number().optional().describe('可选，更新数据库端口号。'),
        username: z.string().optional().describe('可选，更新数据库用户名。'),
        password: z.string().optional().describe('可选，更新数据库密码。'),
        database: z.string().optional().describe('可选，更新数据库名称。'),
        storage: z.string().optional().describe('可选，更新SQLite数据库文件路径。仅当dialect为sqlite时使用。')
      })
    }, async (state: McpToolState) => {
      try {
        const { name, ...updateParams } = state.input;
        
        // 构建更新参数对象
        const updateInfo: Partial<DatabaseConnectionInfo> = {
          ...updateParams
        };
        
        // 如果是sqlite且更新了dialect，设置默认storage
        if (updateParams.dialect === 'sqlite' && !updateParams.storage) {
          updateInfo.storage = path.join(this.dbManager.getDataWiseDir(), `${name}.db`);
        }
        
        const success = await this.dbManager.updateDatabaseConnection(name, updateInfo);
        
        const response = {
          success: true,
          data: {
            name,
            updated: success,
            ...updateParams
          },
          message: success 
            ? `数据库连接 '${name}' 已成功更新并重新连接` 
            : `数据库连接 '${name}' 信息已更新但重新连接失败，请检查连接参数`
        };

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(response, null, 2)
            }
          ]
        };
      } catch (error: any) {
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({
                success: false,
                error: error.message,
                details: '更新数据库连接失败。请检查连接名称是否正确，或者参数是否有效。'
              }, null, 2)
            }
          ]
        };
      }
    });
  }

  async initialize() {
    // 初始化数据库管理器
    try {
      await this.dbManager.initialize();
      console.log('Database manager initialized successfully');
    } catch (error: any) {
      console.error('Failed to initialize database manager:', error.message);
      throw error;
    }
  }

  async start() {
    // 初始化组件
    await this.initialize();

    // 创建传输层
    this.transport = new StdioServerTransport();

    // 启动服务器
    await this.server.connect(this.transport);
    console.log('Database MCP Server is running...');
  }

  async stop() {
    try {
      if (this.transport) {
        await this.transport.close();
      }
      
      // 关闭所有数据库连接
      await this.dbManager.closeAll();
      
      console.log('Database MCP Server stopped.');
    } catch (error: any) {
      console.error('Error while stopping server:', error.message);
    }
  }

  getRelationalDatabaseManager(): DatabaseManager {
    return this.dbManager;
  }
} 